import datetime
import os.path

import numpy as np
from osgeo import gdal, osr


def smap2tif(smap_path, output_dir):
    datasets = gdal.Open(smap_path)
    SubDatasets = datasets.GetSubDatasets()
    # for SubDataset in SubDatasets:
    #     print(SubDataset)
    dataset_am = gdal.Open('HDF5:"{}"://Soil_Moisture_Retrieval_Data_AM/soil_moisture_dca'.format(smap_path))
    dataset_pm = gdal.Open('HDF5:"{}"://Soil_Moisture_Retrieval_Data_PM/soil_moisture_dca_pm'.format(smap_path))

    nodata = dataset_am.GetRasterBand(1).GetNoDataValue()

    data_am = dataset_am.ReadAsArray()
    data_pm = dataset_pm.ReadAsArray()
    data_am[data_am == nodata] = 0
    data_pm[data_pm == nodata] = 0
    data_mean = (data_am + data_pm) / 2
    # print(data_mean)
    data_mean[data_mean == 0] = -9999

    # lonMin = -180
    # lonMax = 180
    # latMin = -85.044
    # latMax = 85.044
    dataset_lon = gdal.Open('HDF5:"{}"://Soil_Moisture_Retrieval_Data_AM/longitude'.format(smap_path))
    dataset_lat = gdal.Open('HDF5:"{}"://Soil_Moisture_Retrieval_Data_AM/latitude'.format(smap_path))
    lon_data = dataset_lon.ReadAsArray()
    lat_data = dataset_lat.ReadAsArray()
    lon_data[lon_data == -9999] = np.nan
    lat_data[lat_data == -9999] = np.nan
    lonMin, lonMax, latMin, latMax = np.nanmin(lon_data), np.nanmax(lon_data), np.nanmin(lat_data), np.nanmax(lat_data)
    # print(np.nanmin(lon_data), np.nanmax(lon_data))
    # print(np.nanmin(lat_data), np.nanmax(lat_data))


    lonRes = (lonMax - lonMin) / dataset_am.RasterXSize
    latRes = (latMax - latMin) / dataset_am.RasterYSize
    # print(lonRes, latRes)

    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    date = os.path.basename(smap_path).split("_")[4]
    targetday = datetime.date(int(str(date)[0:4]), int(str(date)[4:6]), int(str(date)[6:8]))
    day = targetday - datetime.date(targetday.year - 1, 12, 31)
    julian_day = str(date)[0:4] + str(day.days)
    out_tif_path = output_dir + "/" + julian_day + ".tif"
    out_tif = driver.Create(out_tif_path, dataset_am.RasterXSize, dataset_am.RasterYSize, 1, gdal.GDT_Float32)

    # 设置仿射变换矩阵
    geotransform = [lonMin, lonRes, 0, latMax, 0, -latRes]
    out_tif.SetGeoTransform(geotransform)

    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    # 写入数据
    out_tif.GetRasterBand(1).WriteArray(data_mean)
    out_tif.GetRasterBand(1).SetNoDataValue(nodata)
    out_tif.FlushCache()
    out_tif = None
    print("转换完成:{}".format(smap_path))


def main(smap_dir, output_dir):
    smap_list = os.listdir(smap_dir)
    for smap in smap_list:
        smap_path = os.path.join(smap_dir, smap)
        smap2tif(smap_path, output_dir)


if __name__ == '__main__':
    # 单个测试
    # test_path = r"G:\test\SMAP_NEW\SMOS\SM_REPR_MIR_SMUDP2_20160702T202028_20160702T211347_700_200_1"
    # smap2tif(test_path, '')
    # 批量
    smap_dir = r"G:\test\SMAP_NEW\SMAP-36km"
    output_dir = r"G:\test\SMAP_NEW\SMAP-36km-tif"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    main(smap_dir, output_dir)